"""
Post-processing script for HEMC validation.
处理 Geant4 scoring mesh, 要求必须是 1 Dimension 的,
从 .csv 文件读取 quantity 并作图.
"""

import glob
import re
from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

from ..config.models import HEMCConfig
from ..config.router import router

_MODULENAME = Path(__file__).stem


def parse_csv(path: Path) -> dict:
    """
    Parse Geant4 mesh CSV file.

    Args:
        path: Path to CSV file

    Returns:
        Dictionary mapping iZ to (sum_val, sum_val2, n)
    """
    data = {}
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            iz, iphi, ir, sv, sv2, ent = line.split(",")
            iz = int(iz)
            data[iz] = (float(sv), float(sv2), int(ent))
    return data


def calc_mean_var(dat: dict) -> tuple:
    """
    Calculate mean and variance from parsed CSV data.

    Args:
        dat: Dictionary from parse_csv

    Returns:
        Tuple of (iz_arr, mean_arr, var_arr, sigma_arr)
    """
    izs = sorted(dat.keys())
    means, vars_ = [], []
    for iz in izs:
        sv, sv2, n = dat[iz]
        if n == 0:
            m, v = 0.0, 0.0
        else:
            m = sv / n
            v = (sv2 - sv * sv / n) / n  # Population variance
        means.append(m)
        vars_.append(v)
    return np.array(izs), np.array(means), np.array(vars_), np.sqrt(vars_)


def extract_prefix_energy(filename: str) -> tuple:
    """
    Extract prefix and energy from filename.

    Args:
        filename: Name of the CSV file

    Returns:
        Tuple of (prefix, energy) or (None, None) if not found
    """
    pattern = r"(\w+)_(\d+\.?\d*)MeV"
    match = re.match(pattern, filename)

    if match:
        prefix = match.group(1)  # e.g., IDD, LDose, LETd, eDep
        energy = match.group(2)  # e.g., 110.0, 104.0
        return prefix, energy
    else:
        return None, None


def plot_summary(
    prefix: str, file_list: list, config: HEMCConfig, save: bool = True
) -> None:
    """
    Create a summary plot for a given prefix with all energy levels.

    Args:
        prefix: The prefix (e.g., IDD, LDose, LETd, eDep)
        file_list: List of CSV files with the same prefix
        config: HEMC configuration
        save: Whether to save the plot
    """
    plt.figure(figsize=config.plot.figsize)

    # Define a color map for different energies
    # colors = plt.cm.get_cmap('viridis')(np.linspace(0, 1, len(file_list)))
    lenfl = len(file_list)
    lsp = np.linspace(0, 1, lenfl)
    colors = mpl.colormaps["viridis"](lsp)

    for i, f in enumerate(file_list):
        csv_file = Path(f)
        basename = csv_file.stem

        # Extract energy from filename
        _, energy = extract_prefix_energy(basename)

        dat = parse_csv(csv_file)
        iz, mean, var, sigma = calc_mean_var(dat)

        # Plot mean values with different colors
        plt.plot(iz, mean, label=f"{prefix} @ {energy} MeV", color=colors[i], lw=1.2)

        # Optionally also plot the error bands
        plt.fill_between(iz, mean - sigma, mean + sigma, color=colors[i], alpha=0.1)

    # Set labels and title
    ylabel = config.plot.ylabel_dict.get(prefix, prefix)
    plt.xlabel(config.plot.xlabel)
    plt.ylabel(ylabel)
    plt.title(f"{prefix} Summary (All Energies)")
    plt.tight_layout()
    plt.legend(
        bbox_to_anchor=(1.05, 1), loc="upper left"
    )  # Place legend outside plot to avoid overlap

    if save:
        assert config.project_paths.resultdir is not None
        outpng = config.project_paths.resultdir / f"{prefix}_summary.png"
        plt.savefig(outpng, dpi=config.plot.save_dpi, bbox_inches="tight")
        print(f"Saved summary plot -> {outpng}")
    plt.close()


def plot_one(csv_file: Path, config: HEMCConfig, save: bool = True) -> None:
    """
    Process single CSV file and generate plot.

    Args:
        csv_file: Path to input CSV file
        config: HEMC configuration
        save: Whether to save the plot
    """
    basename = csv_file.stem

    # Extract scorer name and energy from filename using regex
    # Pattern matches: prefix_energyMeV (e.g., IDD_110.00MeV, LDose_104.00MeV)
    pattern = r"(\w+)_(\d+\.?\d*)MeV"
    match = re.match(pattern, basename)

    if match:
        scorer = match.group(1)  # e.g., IDD, LDose, LETd, eDep
        energy = match.group(2)  # e.g., 110.00, 104.00
    else:
        scorer = "quantity"
        energy = "unknown"

    ylabel = config.plot.ylabel_dict.get(scorer, scorer)

    dat = parse_csv(csv_file)
    iz, mean, var, sigma = calc_mean_var(dat)

    plt.figure(figsize=config.plot.figsize)
    plt.fill_between(
        iz, mean - sigma, mean + sigma, color="C0", alpha=0.25, label="±1σ"
    )
    plt.plot(iz, mean, lw=1.2, color="C0", label="Mean")
    plt.xlabel(config.plot.xlabel)
    plt.ylabel(ylabel)
    plt.title(f"{scorer} @ {energy} MeV")
    plt.tight_layout()
    plt.legend()

    if save:
        assert config.project_paths.resultdir is not None
        outpng = config.project_paths.resultdir / f"{basename}.png"
        plt.savefig(outpng, dpi=config.plot.save_dpi)
        print(f"Saved  -> {outpng}")
    plt.close()


@router.register()
def plotcsv(config: HEMCConfig):
    """Main entry point for post-processing."""

    # Use whitelist for specific prefixes instead of processing all CSV files
    prefixes = ["eDep", "IDD", "LDose", "LETd"]
    all_files = []

    assert config.project_paths.g4outdir is not None
    for prefix in prefixes:
        pattern = str(config.project_paths.g4outdir / f"{prefix}_*.csv")
        files = glob.glob(pattern)
        all_files.extend(files)

    if not all_files:
        print("No CSV files found matching the prefixes: eDep, IDD, LDose, LETd!")
        return

    # Group files by prefix for summary plots
    prefix_groups = {p: [] for p in prefixes}
    for f in all_files:
        filename = Path(f).stem
        # Extract prefix directly from the beginning of the filename
        for prefix in prefixes:
            if filename.startswith(prefix + "_"):
                prefix_groups[prefix].append(f)
                break

    # Create individual plots for each file
    print("Creating individual plots...")
    for f in all_files:
        print(f"Processing {f} ...")
        plot_one(Path(f), config, save=True)

    # Create summary plots for each prefix with all energy levels
    if False:
        print("Creating summary plots...")
        for prefix in prefixes:
            file_list = prefix_groups[prefix]
            if file_list:  # Only create summary plot if there are files for this prefix
                print(
                    f"Creating summary plot for {prefix} with {len(file_list)} energy levels..."
                )
                plot_summary(prefix, file_list, config, save=True)


if __name__ == "__main__":
    pass
